diff --git a/yarn/src/org/apache/hadoop/yarn/rabit/ApplicationMaster.java b/yarn/src/org/apache/hadoop/yarn/rabit/ApplicationMaster.java
index f51502018..46787ba09 100644
--- a/yarn/src/org/apache/hadoop/yarn/rabit/ApplicationMaster.java
+++ b/yarn/src/org/apache/hadoop/yarn/rabit/ApplicationMaster.java
@@ -216,12 +216,14 @@ public class ApplicationMaster {
             assert (killedTasks.size() + finishedTasks.size() == numTasks);
             success = finishedTasks.size() == numTasks;
             LOG.info("Application completed. Stopping running containers");
-            if (success) {
-              nmClient.stop();
-            }
             diagnostics = "Diagnostics." + ", num_tasks" + this.numTasks
                     + ", finished=" + this.finishedTasks.size() + ", failed="
                     + this.killedTasks.size() + "\n" + this.abortDiagnosis;
+            if (success) {
+              nmClient.stop();
+            } else {
+              throw new Exception("Application not successful" + diagnostics);
+            }
             LOG.info(diagnostics);
         } catch (Exception e) {
             diagnostics = e.toString();
diff --git a/yarn/src/org/apache/hadoop/yarn/rabit/Client.java b/yarn/src/org/apache/hadoop/yarn/rabit/Client.java
index 57d3e9889..63f92d8be 100644
--- a/yarn/src/org/apache/hadoop/yarn/rabit/Client.java
+++ b/yarn/src/org/apache/hadoop/yarn/rabit/Client.java
@@ -198,7 +198,7 @@ public class Client {
         amContainer.setLocalResources(this.setupCacheFiles(appId));
         amContainer.setEnvironment(this.getEnvironment());
         String cmd = "$JAVA_HOME/bin/java"
-                + " -Xmx256M"
+                + " -Xmx900M"
                 + " org.apache.hadoop.yarn.rabit.ApplicationMaster"
                 + this.cacheFileArg + ' ' + this.appArgs + " 1>"
                 + ApplicationConstants.LOG_DIR_EXPANSION_VAR + "/stdout"
@@ -208,7 +208,7 @@ public class Client {
 
         // Set up resource type requirements for ApplicationMaster
         Resource capability = Records.newRecord(Resource.class);
-        capability.setMemory(256);
+        capability.setMemory(1024);
         capability.setVirtualCores(1);
         LOG.info("jobname=" + this.jobName + ",username=" + this.userName);